import torch.nn as nn
import torch

from Models.Backbone import BackboneNet, Convention


class YOLO(BackboneNet):

    def __init__(self, B=2, classes_num=20):
        super(YOLO, self).__init__()
        self.B = B
        self.classes_num = classes_num


        self.Conv_Back = nn.Sequential(
            Convention(1024, 1024, 3, 1, 1, need_bn=False),
            Convention(1024, 1024, 3, 2, 1, need_bn=False),
            Convention(1024, 1024, 3, 1, 1, need_bn=False),
            Convention(1024, 1024, 3, 1, 1, need_bn=False),
        )


        self.Fc = nn.Sequential(
            nn.Linear(7 * 7 * 1024, 4096),
            nn.LeakyReLU(inplace=True, negative_slope=1e-1),
            nn.Linear(4096, 7 * 7 * (B * 5 + classes_num)),
            nn.Sigmoid()
        )
        self.sigmoid = nn.Sigmoid()
        """
        batchx7x7x30让最后一个维度对应的类别为概率和为1 
        """
        # self.softmax = nn.Softmax(dim=3)
    def forward(self, x):
        x = self.Conv_Feature(x)
        x = self.Conv_Semanteme(x)
        x = self.Conv_Back(x)
        x = x.permute(0, 2, 3, 1)
        x = torch.flatten(x, start_dim=1, end_dim=3)
        x = self.Fc(x)
        x = x.view(-1,7,7,(self.B*5 + self.classes_num))
        # x[:,:,:, 0 : self.B * 5] = self.sigmoid(x[:,:,:, 0 : self.B * 5])
        # x[:,:,:, self.B * 5 : ] = self.softmax(x[:,:,:, self.B * 5 : ])
        """
        在pytorch当中注释部分的操作属于inplace操作，而且在官方文档当中，明确表明
        在多交叉熵当中，pytorch不需要使用softmax，因为在计算的时候是包括了这部分的操作的
        并且在yolov1的损失函数当中，计算的类别损失也不是交叉熵
        """
        x = self.sigmoid(x)
        return x

    def initialize_weights(self, net_param_dict):
        for name, m in self.named_modules():
            if isinstance(m, nn.Conv2d):
                torch.nn.init.kaiming_normal_(m.weight.data)
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                torch.nn.init.kaiming_normal_(m.weight.data)
                m.bias.data.zero_()
            elif isinstance(m, Convention):
                m.weight_init()

        self_param_dict = self.state_dict()
        for name, layer in self.named_parameters():
            if name in net_param_dict:
                self_param_dict[name] = net_param_dict[name]
        self.load_state_dict(self_param_dict)
